# coding=utf-8
#*****************************************************************************
#       Copyright (C) 2008 Arnaud Bergeron <abergeron@gmail.com>,
#                          Sébastien Labbé <slabqc@gmail.com>
#
#  Distributed under the terms of the GNU General Public License version 2 (GPLv2)
#
#  The full text of the GPLv2 is available at:
#
#                  http://www.gnu.org/licenses/
#*****************************************************************************
from itertools import count
from weakref import WeakValueDictionary
from sage.structure.sage_object import SageObject
from sage.misc.latex import latex
from sage.rings.integer import Integer
from sage.words.alpha import Alphabet
from sage.words.word_content import is_Content, MakeContent
from sage.words.utils import *

def is_Words(obj):
    r"""
    Returns True if obj is a word set and False otherwise.
    
    EXAMPLES:
        sage: is_Words(33)
        False
        sage: is_Words(Words('ab'))
        True
    """
    return isinstance(obj, Words_g)

def Words(alpha, finite=True, infinite=True, fclass=None, iclass=None):
    r"""
    Defines the set of words over an alphabet.
    
    If you want to build special words, you can use the additional arguments
    fclass and iclass
    
    INPUT:
        alpha -- either an instance of Alphabet or an argument suitable to 
                 build one.
        finite -- boolean (default: True), whether this instance contains 
                  finite words
        infinite -- boolean (default: True), whether this instance contains
                    infinite words
        fclass -- class (default: FiniteWord), the class used for finite words
	    iclass -- class (default: InfiniteWord), the class used for infinite
                  words
    OUTPUT:
        words -- the set of words

    TESTS:
        sage: wab = Words('ab')
        sage: wab == loads(dumps(wab))
        True
        sage: wab is Words('ab')
        True
    """
    from sage.words.word import FiniteWord, InfiniteWord
    if fclass is None:
        fclass = FiniteWord
    if iclass is None:
        iclass = InfiniteWord
    try:
        alpha = Alphabet(alpha)
    except (TypeError, ValueError):
        raise TypeError, "alpha is not an alphabet and cannot be used to build one"
    fclass = fclass if finite else None
    iclass = iclass if infinite else None
    if (alpha, fclass, iclass) not in _words_cache:
        if finite:
            if infinite:
                res = Words_fi(alpha, fclass, iclass)
            else:
                res = Words_f(alpha, fclass)
        else:
            if infinite:
                res = Words_i(alpha, iclass)
            else:
                raise ValueError, "at least one of finite or infinite words must be allowed"
        _words_cache[(alpha, fclass, iclass)] = res
    return _words_cache[(alpha, fclass, iclass)]

_words_cache = WeakValueDictionary()

class Words_g(SageObject):
    def __init__(self, alpha):
        r"""
        TESTS:
            sage: W = sage.words.words.Words_g(Alphabet('ab'))
        """
        self._alpha = alpha
    
    def __reduce__(self):
        r"""
        Support for pickle.
        
        TESTS:
            sage: sage.words.words.Words_g(Alphabet('ab')).__reduce__()
            (<function Words at 0x...>, (Alphabet: ['a', 'b'],))
        """
        return (Words, (self._alpha,) + self._extra_words_args())
    
    def _extra_words_args(self):
        r"""
        Returns any extra arguments used for pickling
        
        EXAMPLES:
            sage: sage.words.words.Words_g(Alphabet('ab'))._extra_words_args()
            ()
        """
        return ()
    
    def alphabet(self):
        r"""
        Returns the alphabet over which self is defined.
        
        EXAMPLES:
            sage: Words('abc').alphabet()
            Alphabet: ['a', 'b', 'c']
        """
        return self._alpha
        
    def dimension(self):
        r"""
        Returns the size of the alphabet.
        
        EXAMPLES:
            sage: Words('abcdef').dimension()
            6
            sage: Words('').dimension()
            0
        """
        return len(self.alphabet())
    
    def __call__(self, obj=None, part=slice(None), format=None):
        r"""
        TESTS:
            sage: from itertools import repeat, count, imap
            sage: W = Words('ab')
            sage: W(format='empty')
            word: 
            sage: W(['a', 'b'], format='empty')
            Traceback (most recent call last):
              ...
            TypeError: trying to build an empty word with something other than None
            sage: W(['a', 'b', 'b'])
            word: abb
            sage: W(['a', 'b', 'b'], format='list')
            word: abb
            sage: W(10, format='list')
            Traceback (most recent call last):
              ...
            TypeError: trying to build a word backed by a list with a sequence not providing the required operations
            sage: Words([0, 1])(lambda x: x%2)
            Infinite word over [0, 1]
            sage: W = Words([0, 1, 2])
            sage: W(lambda x: x%3, slice(0,10))
            word: 0120120120
            sage: W(repeat(1))
            Infinite word over [0, 1, 2]
            sage: Words(xrange(10))(count(), slice(10))
            word: 0123456789
            sage: Words(xrange(10))(count(), slice(10, 0, -2))
            word: 97531
            sage: W(1)
            word: 1
            sage: W('a', format='symbol')
            Traceback (most recent call last):
            ...
            ValueError: symbol ('a') not in alphabet
            sage: W([0, 1])
            word: 01
        """
        from sage.words.word import is_Word
        if format is None:
            if is_Content(obj):
                format = 'content'
            elif obj in self.alphabet():
                format = 'symbol'
        
        if format == 'symbol':
            if obj not in self.alphabet():
                raise ValueError, "symbol (%r) not in alphabet" % obj
            obj = (obj,)
            format = 'list'
        
        if not slice_ok(part):
            raise TypeError, "part is not a slice or has wrong types for elements"
        if format == 'content':
            if not is_Content(obj):
                raise TypeError, "trying to build a word based on raw content with a non-content"
            content = obj
        elif is_Word(obj) and obj.parent() is self:
            return obj[part]
        else:
            content = MakeContent(obj, self.alphabet().to_int, format=format, part=part)
        return self._make_word(content)
    
    def _make_word(self, content):
        r"""
        Method that __call__ delegates to do the actual building of a word 
        from a content object.
        
        TESTS:
            sage: W = sage.words.words.Words_g(Alphabet('ab'))._make_word('content')
            Traceback (most recent call last):
            ...
            NotImplementedError
        """
        raise NotImplementedError
    
    def __contains__(self, w):
        r"""
        Returns True if w is included in self, and False otherwise.
        
        INPUT:
            w -- a word
            
        EXAMPLES:
            sage: Wab = Words('ab')
            sage: Wabc = Words('abc')
            sage: abba = Wab('abba')
            sage: abc = Wabc('abc')
            sage: abba in Wab
            True
            sage: abba in Wabc
            True
            sage: abc in Wabc
            True
            sage: abc in Wab
            False
            sage: Wabc('baab') in Wab # todo: not implemented
            True
        """
        from sage.words.word import is_Word
        if not is_Word(w):
            return False
        if w._parent is self:
            return True
        else:
            return w._parent <= self
    
    def __lt__(self, other):
        r"""
        Returns True if self is a proper subset of other and False otherwise.
        
        TESTS:
            sage: Words('ab') < Words('ab')
            False
            sage: Words('ab') < Words('abc')
            True
            sage: Words('abc') < Words('ab')
            False
            sage: Words('ab', finite=False) < Words('ab')
            True
            sage: Words('ab') < Words('ab', infinite=False)
            False
            sage: Words('ab') < Words('abc', infinite=False)
            False
        """
        if not is_Words(other):
            return NotImplemented
        return self <= other and self != other
    
    def __gt__(self, other):
        r"""
        Returns True if self is a proper superset of other and False otherwise.
        
        TESTS:
            sage: Words('ab') > Words('ab')
            False
            sage: Words('ab') > Words('abc')
            False
            sage: Words('abc') > Words('ab')
            True
            sage: Words('ab', finite=False) > Words('ab')
            False
            sage: Words('ab') > Words('ab', infinite=False)
            True
            sage: Words('abc', infinite=False) > Words('ab')
            False
        """
        if not is_Words(other):
            return NotImplemented
        return self >= other and self != other
    
    def __le__(self, other):
        r"""
        Returns True if self is a subset of other and False otherwise.
        
        TESTS:
            sage: Words('ab') <= Words('ab')
            True
            sage: Words('ab') <= Words('abc')
            True
            sage: Words('abc') <= Words('ab')
            False
            sage: Words('ab', finite=False) <= Words('ab')
            True
            sage: Words('ab') <= Words('ab', infinite=False)
            False
            sage: Words('ab') <= Words('abc', infinite=False)
            False
        """
        if not is_Words(other):
            return NotImplemented
        if isinstance(other, type(self)):
            return self.alphabet() <= other.alphabet()
        else:
            return False

    def __ge__(self, other):
        r"""
        Returns True if self is a superset of other and False otherwise.
        
        TESTS:
            sage: Words('ab') >= Words('ab')
            True
            sage: Words('ab') >= Words('abc')
            False
            sage: Words('abc') >= Words('ab')
            True
            sage: Words('ab', finite=False) >= Words('ab')
            False
            sage: Words('ab') >= Words('ab', infinite=False)
            True
            sage: Words('abc', infinite=False) >= Words('ab')
            False
        """
        if not is_Words(other):
            return NotImplemented
        if isinstance(self, type(other)):
            return self.alphabet() >= other.alphabet()
        else:
            return False
    
class Words_f(Words_g):
    def __init__(self, alpha, fclass):
        r"""
        TESTS:
            sage: W = Words('ab', infinite=False)
            sage: W is loads(dumps(W))
            True
        """
        Words_g.__init__(self, alpha)
        self._fclass = fclass
    
    def _extra_words_args(self):
        r"""
        Returns any extra arguments used for pickling.
        
        EXAMPLES:
            sage: from sage.words.word import FiniteWord
            sage: Words('ab', infinite=False)._extra_words_args()
            (True, False, <class 'sage.words.word.FiniteWord'>, None)
        """
        return (True, False, self._fclass, None)
    
    def iter_len(self, l=1):
        r"""
        Returns an iterator over all the words of self of length l.
        
        INPUT:
            l -- integer (default: 1) the length of the desired words
        
        EXAMPLES:
            sage: W = Words('ab', infinite=False)
            sage: list(W.iter_len()) 
            [word: a, word: b]
            sage: list(W.iter_len(2))
            [word: aa, word: ba, word: ab, word: bb]
            sage: list(W.iter_len(3))
            [word: aaa,
             word: baa,
             word: aba,
             word: bba,
             word: aab,
             word: bab,
             word: abb,
             word: bbb]
            sage: list(W.iter_len('a'))
            Traceback (most recent call last):
            ...
            TypeError: the parameter l (='a') must be an integer
        """
        if not isint(l):
            raise TypeError, "the parameter l (=%r) must be an integer"%l
            
        i = Integer(0)
        dim = self.dimension()
        alphabet = list(self.alphabet())
        iterations = dim ** l
        while i < iterations:
            yield self(i.digits(base=dim, digits=alphabet, padto=l))
            i+=1
        else:
            raise StopIteration
            
    def __iter__(self):
        r"""
        Returns an iterator over all the words of self.
        
        EXAMPLES:
            sage: W=Words([4,5], infinite=False)
            sage: for w in W:
            ...     if len(w)>3:
            ...         break
            ...     else:
            ...         print w
            ...         
            word: 
            word: 4
            word: 5
            word: 44
            word: 54
            word: 45
            word: 55
            word: 444
            word: 544
            word: 454
            word: 554
            word: 445
            word: 545
            word: 455
            word: 555
        """
        for l in count():
            for w in self.iter_len(l):
                yield w
        
    def iter_morphisms(self, l):
        r"""
        Returns an iterator over all the endomorphisms $\varphi$ on self 
        such that $|\varphi(a)|=l[i]$ for all i-th letters $a$ of self (in the
        order defined on the alphabet).
                
        INPUT:
            l -- list of integers such that len(l) == self.dimension()
            
        OUTPUT:
            iterator -- over the morphisms
            
        EXAMPLES:
            sage: W = Words('ab', infinite=False)                 
            sage: map(str, W.iter_morphisms([2, 1]))
            ['WordMorphism: a->aa, b->a',
             'WordMorphism: a->ba, b->a',
             'WordMorphism: a->ab, b->a',
             'WordMorphism: a->bb, b->a',
             'WordMorphism: a->aa, b->b',
             'WordMorphism: a->ba, b->b',
             'WordMorphism: a->ab, b->b',
             'WordMorphism: a->bb, b->b']
            sage: map(str, W.iter_morphisms([2, 2]))
            ['WordMorphism: a->aa, b->aa',
             'WordMorphism: a->ba, b->aa',
             'WordMorphism: a->ab, b->aa',
             'WordMorphism: a->bb, b->aa',
             'WordMorphism: a->aa, b->ba',
             'WordMorphism: a->ba, b->ba',
             'WordMorphism: a->ab, b->ba',
             'WordMorphism: a->bb, b->ba',
             'WordMorphism: a->aa, b->ab',
             'WordMorphism: a->ba, b->ab',
             'WordMorphism: a->ab, b->ab',
             'WordMorphism: a->bb, b->ab',
             'WordMorphism: a->aa, b->bb',
             'WordMorphism: a->ba, b->bb',
             'WordMorphism: a->ab, b->bb',
             'WordMorphism: a->bb, b->bb']
            sage: map(str, W.iter_morphisms([0, 0]))
            ['WordMorphism: a->, b->']
            sage: map(str, W.iter_morphisms([0, 1]))
            ['WordMorphism: a->, b->a', 'WordMorphism: a->, b->b']
            
        TESTS:
            sage: list(W.iter_morphisms([0, 1, 2]))
            Traceback (most recent call last):
            ...
            TypeError: l (=[0, 1, 2]) must be a list of 2 integers
            sage: list(W.iter_morphisms([0, 'a'])) 
            Traceback (most recent call last):
            ...
            TypeError: l (=[0, 'a']) must be a list of 2 integers
        """
        if not isinstance(l, list) or not len(l) == self.dimension() \
            or not all(map(isint,l)):
            raise TypeError, "l (=%s) must be a list of %s integers" \
                             %(l, self.dimension())

        from sage.words.morphism import WordMorphism
        
        cuts = [0] + l
        for i in range(1,len(cuts)):
            cuts[i] += cuts[i-1]
        
        s = cuts[-1] # same but better than s = sum(l)
        for big_word in self.iter_len(s):
            d = {}
            i = 0
            for a in self.alphabet():
                d[a] = big_word[cuts[i]:cuts[i+1]]
                i += 1
            yield WordMorphism(d)
    
    def _repr_(self):
        r"""
        Returns a string representation of self.
        
        EXAMPLES:
            sage: Words('ab', infinite=False)._repr_()
            "Finite words over ['a', 'b']"
        """
        return "Finite words over %s" % self._alpha.string_rep()
    
    def _latex_(self):
        r"""
        Returns a \LaTeX representation of the object.
        
        TESTS:
            sage: Words('ab', infinite=False)._latex_()
            '[\\text{a}, \\text{b}]^*'
        """
        return latex(self._alpha) + "^*"
    
    def _make_word(self, content):
        r"""
        TESTS:
            sage: from sage.words.word_content import MakeContent
            sage: w = Words('ab', infinite=False)._make_word(MakeContent([0, 1, 1, 0]))
            sage: w = Words('ab', infinite=False)._make_word(MakeContent(lambda x: x%2))
            Traceback (most recent call last):
            ...
            TypeError: infinite words are unsupported for this set
        """
        if not haslen(content):
            raise TypeError, "infinite words are unsupported for this set"
        return self._fclass(content, self)

class Words_i(Words_g):
    def __init__(self, alpha, iclass):
        r"""
        TESTS:
            sage: W = Words('ab', finite=False)
            sage: W is loads(dumps(W))
            True
        """
        Words_g.__init__(self, alpha)
        self._iclass = iclass
    
    def _extra_words_args(self):
        r"""
        Returns any extra arguments used for pickling.
        
        EXAMPLES:
            sage: Words('ab', finite=False)._extra_words_args()
            (False, True, None, <class 'sage.words.word.InfiniteWord'>)
        """
        return (False, True, None, self._iclass)
    
    def _repr_(self):
        r"""
        Returns a string representation of self.
        
        EXAMPLES:
            sage: Words('ab', finite=False)._repr_()
            "Infinite words over ['a', 'b']"
        """
        return "Infinite words over %s" % self._alpha.string_rep()
    
    def _latex_(self):
        r"""
        Returns a \LaTeX representation of the object.
        
        TESTS:
            sage: Words('ab', finite=False)._latex_()
            '[\\text{a}, \\text{b}]^\\omega'
        """
        return latex(self._alpha) + "^\\omega"
    
    def _make_word(self, content):
        r"""
        TESTS:
            sage: from sage.words.word_content import MakeContent
            sage: from itertools import cycle
            sage: w = Words('ab', finite=False)._make_word(MakeContent(cycle([0, 1, 1, 0])))
            sage: w = Words('ab', finite=False)._make_word(MakeContent([0, 1, 1, 0]))
            Traceback (most recent call last):
            ...
            TypeError: finite words are unsupported for this set
        """
        if haslen(content):
            raise TypeError, "finite words are unsupported for this set"
        return self._iclass(content, self)

class Words_fi(Words_f, Words_i):
    def __init__(self, alpha, fclass, iclass):
        r"""
        TESTS:
            sage: W = Words('ab')
            sage: W is loads(dumps(W))
            True
        """
        Words_f.__init__(self, alpha, fclass)
        Words_i.__init__(self, alpha, iclass)
    
    def _extra_words_args(self):
        r"""
        Returns any extra arguments used for pickling.
        
        EXAMPLES:
            sage: Words('ab')._extra_words_args()
            (True, True, <class 'sage.words.word.FiniteWord'>, <class 'sage.words.word.InfiniteWord'>)
        """
        return (True, True, self._fclass, self._iclass)

    
    def _repr_(self):
        r"""
        Returns a string representation of self.
        
        TESTS:
            sage: Words('ab')._repr_()
            "Words over ['a', 'b']"
        """
        return "Words over %s" % self._alpha.string_rep()
    
    def _latex_(self):
        r"""
        Returns a \LaTeX representation of the object.
        
        TESTS:
            sage: Words('ab')._latex_()
            '[\\text{a}, \\text{b}]^* \\cup [\\text{a}, \\text{b}]^\\omega'
        """
        return Words_f._latex_(self) + " \\cup " + Words_i._latex_(self)
    
    def _make_word(self, content):
        r"""
        TESTS:
            sage: from sage.words.word_content import MakeContent
            sage: from itertools import cycle
            sage: w = Words('ab')._make_word(MakeContent([0, 1, 1, 0]))
            sage: w = Words('ab')._make_word(MakeContent(cycle([0, 1, 1, 0])))
        """
        if haslen(content):
            return Words_f._make_word(self, content)
        else:
            return Words_i._make_word(self, content)
